import platform
import pickle,asyncio
 
if platform.system() == "Windows":
    import asyncio 
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
import datetime
import os
import warnings
from playwright.async_api import async_playwright

import matplotlib.pyplot as plt
import matplotlib
import pandas as pd
import seaborn as sns
from IPython.display import display

warnings.filterwarnings("ignore")
matplotlib.rcParams['axes.unicode_minus']=False#显示负号
plt.rcParams['font.sans-serif'] = ['simhei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

##################################

      
class Stock:
    def __init__(self,n,days):
        '初始化股票数据'
        self.days=days
        self.n=n
        self.loadData()
        
    async def getName(self):
        '使用雪球api获取股票名称'
        async with async_playwright() as p:            
            browser = await p.chromium.launch(headless=False,slow_mo=2000)
            page = await browser.new_page()
            for code in self.stocks:  
                name=self.names.get(code,None)
                if name==None:
                    await page.goto(f'https://xueqiu.com/S/{code}')
                    name=await page.query_selector('.stock-name')
                    name=await name.text_content()
                    self.names[code]=name
                self.stocks[code].append(name)
                path=os.getcwd()
                namePath=os.path.join(path,'nams.dat')
                with open(namePath,'wb') as f:
                    pickle.dump(self.names,f)
            await browser.close()

    
    def loadData(self):
        '加载股票数据'
        self.stocks={}
        path=os.getcwd()
        namePath=os.path.join(path,'nams.dat')
        if os.path.isfile(namePath):
            with open(namePath,'rb') as f:
                self.names=pickle.load(f)
        else:
            self.names={}            
        path=os.path.join(path,'stocks')
        filesName=os.listdir(path)
        for filename in filesName:
            if '.' in filename:
                code=filename.split('.')[0][-8:]
                self.stocks[code]=[os.path.join(path,filename)]


        
        
    def getData(self):
        '拉取数据'        
        code=self.code
        name=self.name
        print(f'开始拉取【{name}】的数据...')    
        datas=pd.read_csv(self.stocks[code][0],index_col='日期 - date')
        datas=datas[:1001]
        datas=datas[::-1]
        datas = datas[
            [
                '滚动市盈率 - pe_ttm',
                '滚动扣非市盈率 - koufei_pe_ttm',
                '滚动市净率 - pb_ttm',
                '动态市销率 - ps_ttm',
                r'动态股息收益率 % - dividend_yield',
                'ROE %',
                '总市值 - total market cap',
                'A股总市值 - A gu zong shi zhi',
                '收盘价(前复权) - close price',
                '开盘价-open',
                '涨跌幅-up',
            ]
        ]  
        up=list(datas['涨跌幅-up'])
        datas['涨跌幅-up']=[float(i[:-1])/100 for i in up]   
        self.df=datas

        
    def bsTest(self,b,s):
        '根据bs点买卖的回测'       
        isb=0
        startMoney=1000000
        money=startMoney
        bs={
            '买入日期':[],
            '买入价格':[],
            f'买入时{self.factor}':[],
            '卖出日期':[],
            '卖出价格':[],
            f'卖出时{self.factor}':[],
            '本次盈利':[],
            '持有天数':[],
        }
        for index in self.df.index:            
            data=self.df[self.factor][index]
            price=self.df['收盘价(前复权) - close price'][index]
            #正相关低买高卖，负相关低卖高买 
            def buy(isb,money):
                '买入'
                isb=1
                num=money//(100*price)
                left=money-num*price*100
                bDate=index
                bPrice=price
                bData=data
                #print(bDate,'买入',end='\r')
                return isb,num,left,bDate,bPrice,bData
            def sall(isb,num,left,bDate,bPrice,bData):
                '卖出'
                isb=0
                money=num*price*100+left
                sDate=index
                sPrice=price
                bs['买入日期'].append(bDate)
                bs['买入价格'].append(bPrice)
                bs[f'买入时{self.factor}'].append(bData)
                bs['卖出日期'].append(sDate)
                bs['卖出价格'].append(sPrice)
                bs[f'卖出时{self.factor}'].append(data)
                bs['本次盈利'].append(sPrice/bPrice-1)
                sDate=datetime.datetime.strptime(sDate, '%Y-%m-%d')
                bDate=datetime.datetime.strptime(bDate, '%Y-%m-%d')
                bs['持有天数'].append(sDate-bDate)
                #print(sDate,'卖出',end='\r')
                return isb,money,sDate,sPrice,bs
            if self.corr>0:
                if isb==0 and data<b:
                    isb,num,left,bDate,bPrice,bData=buy(isb,money)                        
                elif isb==1 and data>=s:                     
                    isb,money,sDate,sPrice,bs=sall(isb,num,left,bDate,bPrice,bData)
                    
            else:
                if isb==0 and data>b:
                    isb,num,left,bDate,bPrice,bData=buy(isb,money)
                elif isb==1 and data<=s:
                    isb,money,sDate,sPrice,bs=sall(isb,num,left,bDate,bPrice,bData)
        if isb:
            isb,money,sDate,sPrice,bs=sall(isb,num,left,bDate,bPrice,bData)
            hold=1
        else:
            hold=0
        bs=pd.DataFrame(bs)
        return money/startMoney,bs,hold
    def testFactor(self):
        '研究单个因子'  
        n=self.n
        datas=self.df[self.factor]
        step=(max(datas)-min(datas))/n
        bPoint=[min(datas)+i*step for i in range(n+1)]
        sPoint=[min(datas)+i*step for i in range(n+1)]
        result={
            f'买入{self.factor}':[],
            f'卖出{self.factor}':[],
            '最终净值':[],
            '交易次数':[],
            '胜率':[],
            '终日是否持有':[],
        }
        for b in bPoint:
            if self.corr<0:
                s=sPoint[0] 
            else:
                s=sPoint[-1] 
            trueS=s
            maxIncome,bs,isHold=self.bsTest(b,s)    
            for s in sPoint:
                if self.corr>0:
                    if b<s:                        
                        now,nowBS,hold=self.bsTest(b,s)
                        if now >maxIncome:
                            maxIncome,bs,isHold=now,nowBS,hold
                            trueS=s
                else:
                    if b>s:
                        now,nowBS,hold=self.bsTest(b,s)
                        if now >maxIncome:
                            maxIncome,bs,isHold=now,nowBS,hold
                            trueS=s                
            c=0
            for i in bs['本次盈利']:
                if i>0:
                    c+=1            
            if len(bs)>0:
                ratio=c/len(bs)            
                result[f'买入{self.factor}'].append(b)
                result[f'卖出{self.factor}'].append(trueS)
                result['最终净值'].append(maxIncome)
                result['终日是否持有'].append(isHold)
                result['胜率'].append(ratio)
                result['交易次数'].append(len(bs))
        df=pd.DataFrame(result)        
        df=df.loc[df["交易次数"] >=5]
        df=df.sort_values(by='最终净值')
        n=len(df)
        if n>15:
            df=df[-1:-16:-1]
        else:
            df=df[::-1]
        display(df)
        
    def run(self):
        '研究单支股票'        
        self.getData()
        print(f'【{self.name}】各种因子相关性分析：')
        df=self.df.corr(method='spearman')        
        display(df)
        
        #sns.clustermap(df)   
        sns.pairplot(
            self.df,
            #kind='reg',
            diag_kind='kde',
            x_vars=[i for i in self.df.columns if not i =='收盘价(前复权) - close price'],
            y_vars=['收盘价(前复权) - close price'])
        plt.title(f'【{self.name}】各种因子相关性分析',loc='center')
        plt.show()
        
        aFactor=[i for i in df.columns[1:]]
        self.result={}        
        maxFN=aFactor[0]
        maxFD=df['收盘价(前复权) - close price'][maxFN]
        for factor in aFactor[:-1]:
            corr=df['收盘价(前复权) - close price'][factor]
            if abs(corr)>abs(maxFD) and factor not in ['开盘价-open','收盘价(前复权) - close price']:
                maxFN,maxFD=factor,corr
        print(f'【{self.name}】经过分析，与收盘价相关性最高的因子是【{maxFN}】，分析数据：')
        for factor in aFactor:
            corr=df['收盘价(前复权) - close price'][factor]  
            if abs(corr) <=0.2:
                print(f'    【{factor}】因子的相关性太低了，跳过分析！')
            else:
                print(f'【{self.name}】【{factor}】因子的相关性为【{corr}】，回测详情如下：')
                self.factor=factor
                self.corr=corr
                self.result=self.testFactor()     
            if corr>maxFD:
                maxFN,maxFD=factor,corr
        print('近期数据如下：')
        display(self.df.tail())
        
    def main(self):
        '主程序'
        while True:
            print('检索到股票数据如下：')
            count=0
            for code in self.stocks:
                name=self.stocks[code][-1]
                self.name=name
                self.code=code
                print(f'{count}：{name}')
                count+=1
            print('-'*50)
            n=input('请输入要查询的股票序号：\n>>> ')
            n=int(n)
            code=list(self.stocks)[n]
            name=self.stocks[code][-1]
            self.name=name
            self.code=code
            self.run()
            print('#'*50)

s=Stock(100,1000)
asyncio.run(s.getName())
s.main()